import os
import pandas as pd
from pathlib import Path
import dolphindb as ddb
import numpy as np
import traceback

class SHFEImporter:
    """上海期货交易所数据导入器"""
    
    def __init__(self):
        self.source_path = Path("F:/资料/历史数据/上期所")
        self.session = None
        self.upserter = None
        
    def connect_db(self):
        """连接DolphinDB数据库"""
        self.session = ddb.session()
        self.session.connect("localhost", 8848, "admin", "123456")
        self.upserter = ddb.TableUpserter(
            dbPath="dfs://futures", 
            tableName="daily_all", 
            ddbSession=self.session,
            keyColNames=["date", "code", "exchange_code"]
        )
        
    def _read_excel(self, file_path):
        """使用多种方法尝试读取Excel文件"""
        engines = {
            'xlrd': {'header': 2},
            'openpyxl': {'header': 2},
            'pyxlsb': {'header': 2}
        }
        
        for engine, params in engines.items():
            try:
                df = pd.read_excel(file_path, engine=engine, **params)
                return df
            except Exception as e:
                continue
                
        return None
            
    def is_option_code(self, code):
        """判断是否为期权合约"""
        try:
            code = str(code).upper()
            option_indicators = ['P', 'C']
            return any(indicator in code for indicator in option_indicators) and len(code) > 6
        except:
            return False

    def _get_product_name(self, code):
        """从合约代码提取品种名称"""
        try:
            code = str(code)
            product_names = {
                'cu': '铜',
                'al': '铝',
                'zn': '锌',
                'pb': '铅',
                'ni': '镍',
                'sn': '锡',
                'au': '黄金',
                'ag': '白银',
                'rb': '螺纹钢',
                'wr': '线材',
                'hc': '热轧卷板',
                'fu': '燃料油',
                'bu': '沥青',
                'ru': '天然橡胶'
            }
            product_code = ''.join(filter(str.isalpha, code[:2].lower()))
            return product_names.get(product_code, product_code)
        except:
            return code

    def process_file(self, file_path, year):
        """处理单个文件"""
        try:
            df = self._read_excel(file_path)
            if df is None:
                return None

            # 根据年份决定数据起始行
            data_start_row = 1 if int(year) >= 2020 else 0
            df = df.iloc[data_start_row:]
            
            # 处理合并单元格
            df['合约'] = df['合约'].ffill()
            
            # 基础数据清理
            df = df.dropna(how='all')
            df = df[~df['合约'].str.contains('注：', na=False)]
            
            # 数值转换
            numeric_cols = {
                '开盘价': 'open', '最高价': 'high', '最低价': 'low',
                '收盘价': 'close', '前收盘': 'pre_close', '成交量': 'volume'
            }
            
            for cn_col, en_col in numeric_cols.items():
                if cn_col in df.columns:
                    df[cn_col] = pd.to_numeric(df[cn_col], errors='coerce')
            
            # 处理成交金额
            amount_cols = ['成交金额', '成交金额(万元)']
            for col in amount_cols:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')
                    if '万元' in col:
                        df[col] *= 10000
                    break
            
            # 重命名列
            rename_dict = {
                '日期': 'date',
                '合约': 'code',
                '开盘价': 'open',
                '最高价': 'high',
                '最低价': 'low',
                '收盘价': 'close',
                '前收盘': 'pre_close',
                '成交量': 'volume',
                '成交金额': 'amount',
                '成交金额(万元)': 'amount'
            }
            
            df = df.rename(columns={k: v for k, v in rename_dict.items() if k in df.columns})
            
            # 处理日期格式
            df['date'] = pd.to_datetime(df['date'], format='%Y%m%d')
            
            # 添加交易所信息
            df['exchange_code'] = 'SHFE'
            df['exchange_name'] = '上期所'
            
            # 提取品种名称并过滤期权
            df['name'] = df['code'].apply(self._get_product_name)
            df = df[~df['code'].apply(self.is_option_code)]
            
            # 选择和排序最终列
            final_columns = ['date', 'code', 'name', 'exchange_code', 'exchange_name',
                           'high', 'low', 'open', 'close', 'pre_close', 'volume', 'amount']
            
            return df[final_columns]
            
        except Exception as e:
            print(f"处理文件失败 {file_path}: {str(e)}")
            print(f"错误类型: {type(e).__name__}")
            print(f"详细错误信息: {traceback.format_exc()}")
            return None

    def import_to_db(self, df):
        """导入数据到数据库"""
        try:
            self.upserter.upsert(df)
            print(f"成功处理 {len(df)} 条数据")
        except Exception as e:
            print(f"导入数据失败: {str(e)}")
            try:
                print("尝试重新连接...")
                self.connect_db()
                self.upserter.upsert(df)
                print(f"重连后成功处理 {len(df)} 条数据")
            except Exception as e2:
                print(f"重试失败: {str(e2)}")

    def process_all_files(self):
        """处理所有文件"""
        try:
            self.connect_db()
            
            for year_dir in self.source_path.iterdir():
                if not year_dir.is_dir():
                    continue
                    
                year = year_dir.name
                print(f"\n处理 {year} 年数据...")
                
                for file_path in year_dir.glob('*.xls*'):
                    print(f"处理文件: {file_path.name}")
                    df = self.process_file(file_path, year)
                    
                    if df is not None and not df.empty:
                        print(f"处理得到 {len(df)} 行数据")
                        self.import_to_db(df)
                        
        except Exception as e:
            print(f"处理过程出错: {str(e)}")
            raise
        finally:
            if self.session:
                self.session.close()

def main():
    """主函数"""
    importer = SHFEImporter()
    importer.process_all_files()

if __name__ == "__main__":
    main() 